package com.jeesite.modules.dp.service;

import com.alibaba.fastjson.JSON;
import com.jeesite.common.lang.StringUtils;
import io.github.imfangs.dify.client.DifyClientFactory;
import io.github.imfangs.dify.client.DifyWorkflowClient;
import io.github.imfangs.dify.client.callback.WorkflowStreamCallback;
import io.github.imfangs.dify.client.enums.ResponseMode;
import io.github.imfangs.dify.client.event.*;
import io.github.imfangs.dify.client.model.workflow.WorkflowRunRequest;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;

import javax.security.auth.message.MessageInfo;
import javax.websocket.*;
import javax.websocket.server.PathParam;
import javax.websocket.server.ServerEndpoint;
import java.io.IOException;
import java.util.*;
import java.util.concurrent.ConcurrentHashMap;

@Slf4j
@Component
@ServerEndpoint("/a/dify/{userId}")
public class WebSocketServer {


    //存储客户端session信息
    public static Map<String, Session> clients = new ConcurrentHashMap<>();

    //存储把不同用户的客户端session信息集合
    public static Map<String, Set<String>> connection = new ConcurrentHashMap<>();

    //会话id
    private String sid = null;

    //建立连接的用户id
    private String userId;

    @OnOpen
    public void onOpen(Session session, @PathParam("userId") String userId) {
        this.sid = UUID.randomUUID().toString();
        this.userId = userId;
        clients.put(this.sid, session);
        //判断该用户是否存在会话信息，不存在则添加
        Set<String> clientSet = connection.computeIfAbsent(userId, k -> new HashSet<>());
        clientSet.add(this.sid);
        log.info("{}用户建立连接，{}连接开启！", this.userId, this.sid);
    }

    @OnClose
    public void onClose() {
        clients.remove(this.sid);
        log.info("{}连接断开", this.sid);
    }

    DifyWorkflowClient chatClient = DifyClientFactory
            .createWorkflowClient("https://api.dify.ai/v1",
                    "app-kCyYtsgoUauFa1oIT21ZmCF7");

    @OnMessage
    public void onMessage(String message, Session session) {
        log.info("收到来自用户：{}的信息   {}", this.userId, message);
        try {

            Map<String, Object> inputs = new HashMap<>();
            inputs.put("input_start", message);

            WorkflowRunRequest request = WorkflowRunRequest.builder()
                    .inputs(inputs)
                    .user("user-123")
                    .responseMode(ResponseMode.STREAMING)
                    .build();
            chatClient.runWorkflowStream(request, new WorkflowStreamCallback() {
                @Override
                public void onWorkflowStarted(WorkflowStartedEvent event) {
                    System.out.println("started:" + event.toString());
                }

                @Override
                public void onNodeStarted(NodeStartedEvent event) {
                }

                @Override
                public void onNodeFinished(NodeFinishedEvent event) {
                }

                @Override
                public void onWorkflowFinished(WorkflowFinishedEvent event) {
                    session.getAsyncRemote().sendText("θ===end===θ");
                }

                @Override
                public void onWorkflowTextChunk(WorkflowTextChunkEvent event) {
                    session.getAsyncRemote().sendText(event.getData().getText());
                }

                @Override
                public void onTtsMessage(TtsMessageEvent event) {
                    session.getAsyncRemote().sendText(event.getEvent());
                }

                @Override
                public void onTtsMessageEnd(TtsMessageEndEvent event) {
                    session.getAsyncRemote().sendText("===message end===");
                }

            });
        } catch (IOException e) {
            session.getAsyncRemote().sendText(e.getMessage());
        }

    }

    @OnError
    public void onError(Throwable error) {
        log.info("系统错误");
        error.printStackTrace();
    }

    /**
     * @description: 通过userId向用户发送信息
     * 该类定义成静态可以配合定时任务实现定时推送
     **/
    public static void sendMessageByUserId(String userId, MessageInfo message) {
        if (!StringUtils.isEmpty(userId)) {
            Set<String> clientSet = connection.get(userId);
            //用户是否存在客户端连接
            if (Objects.nonNull(clientSet)) {
                Iterator<String> iterator = clientSet.iterator();
                while (iterator.hasNext()) {
                    String sid = iterator.next();
                    Session session = clients.get(sid);
                    //向每个会话发送消息
                    if (Objects.nonNull(session)) {
                        try {
                            String jsonString = JSON.toJSONString(message);
                            //同步发送数据，需要等上一个sendText发送完成才执行下一个发送
                            session.getBasicRemote().sendText(jsonString);
                        } catch (IOException e) {
                            e.printStackTrace();
                        }
                    }
                }
            }
        }
    }
}
